package tool.code.generater.sampler;


import tool.code.generater.DataType;
import tool.code.generater.config.JdbcConfig;
import tool.code.generater.struct.ColumnDef;
import tool.code.generater.struct.TableDef;
import tool.code.generater.util.DataSourceHelper;
import com.alibaba.druid.DbType;
import org.apache.commons.lang3.StringUtils;

import java.sql.JDBCType;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

public class DBTableInfoSampler extends TableInfoSampler {
    private final DataSourceHelper dataSourceHelper;

    public DBTableInfoSampler(DbType dbType, String mappingConfigFile, JdbcConfig jdbcConfig) {
        super(dbType, mappingConfigFile, jdbcConfig);
        this.dataSourceHelper = new DataSourceHelper(jdbcConfig);
    }

    @Override
    public List<TableDef> getColumns() {
        String[] tables = this.jdbcConfig.getTables();
        List<TableDef> tableDefs = new ArrayList<TableDef>();
        try {
            Map<String, String> tablesComment = dataSourceHelper.getTablesComment();
            for (String table : tables) {
                List<Map<String, Object>> columns = dataSourceHelper.getColumns(table);
                TableDef tableDef = new TableDef();
                tableDef.setName(table);
                List<ColumnDef> columnDefs = new ArrayList<ColumnDef>();
                for (Map<String, Object> column : columns) {
                    ColumnDef columnDef = new ColumnDef();
                    columnDef.setPrimaryKey(column.containsKey("primaryKey"));
                    columnDef.setName((String) column.get("name"));
                    String dataType = String.valueOf(column.get("dataType"));
                    int type = Integer.parseInt(dataType);
                    DataType fieldType = getFieldType(type);
                    columnDef.setType(dataType);
                    columnDef.setJavaTypePackage(fieldType.getJavaTypePackage());
                    columnDef.setFieldType(fieldType.getJavaType());
                    columnDef.setComment((String) column.get("remarks"));
                    columnDefs.add(columnDef);
                }
                tableDef.setColumns(columnDefs);
                tableDef.setComment(tablesComment.get(table));
                Set<String> collect = columnDefs.stream().map(ColumnDef::getJavaTypePackage)
                        .filter(StringUtils::isNotEmpty)
                        .collect(Collectors.toSet());
                tableDef.setJavaTypePackages(new ArrayList<>(collect));
                tableDefs.add(tableDef);
            }
        } catch (SQLException e) {
            throw new RuntimeException(e);
        } finally {
            dataSourceHelper.close(dataSourceHelper.getDataSource());
        }
        return tableDefs;
    }

    protected DataType getFieldType(int type) {
        if (this.typeMapper == null) {
            return null;
        }
        JDBCType jdbcType = JDBCType.valueOf(type);
        return this.typeMapper.getJavaType(jdbcType.getName());
    }

}
